'''
Created on May 21, 2012

@author: tel
'''
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from pylab import setp
import stats
import sys, copy, os, pickle
import scipy.stats as ss
from plothelper import Unpickle, Repickle, Polyfit

def OffTicks(ax, list):
    for entry in list:
        if entry=='x':
            setp(ax.get_xticklabels(), visible=False)
        elif entry=='y':
            setp(ax.get_yticklabels(), visible=False)

def Scatter(*args, **kwargs):
    if 'outfile' in kwargs:
        outfile = kwargs['outfile']
    else:
        outfile = 'Scatter'
    if 'outdir' in kwargs:
        outdir = kwargs['outdir']
    else:
        outdir = '/home/tel/benchmark/links'
    if 'repickle' in kwargs:
        repickle = kwargs['repickle']
    else:
        repickle = True
    
    if repickle:
        Repickle(outdir, outfile, *args)
        
    if outfile=='dsol_ele_scatter' or outfile=='dsol_ele_csa_scatter':
        xlabel = '$\Delta G_{dsol}$'
        ylabel = '$\Delta G_{ele}$'
    elif outfile=='dsol_ele_csa_hist':
        xlabel = 'ratio'
        ylabel = 'count'
    elif outfile=='ddpk_mddpk_scatter':
        xlabel = '$\Delta\Delta pK_{a}$'
        ylabel = 'model $\Delta\Delta pK_{a}$'
    elif outfile=='ratio_ddpk_scatter':
        xlabel =r'$\frac{\Delta G_{ele}}{\Delta G_{dsol}}$'
        ylabel = '$\Delta\Delta pK_{a}$'
    
    fig = plt.figure(1)
    axes = plt.axes()
    
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)

    for tick in axes.xaxis.get_major_ticks():
        tick.label.set_fontsize('small')
        tick.label.set_rotation(45)
    for tick in axes.yaxis.get_major_ticks():
        tick.label.set_fontsize('small')
    
    if outfile=='dsol_ele_csa_scatter':
        faraway = zip(*args)
        faraway = filter(lambda x:10<x[2]<999, faraway)
        fx,fy = zip(*faraway)[0:2]
        axes.scatter(fx, fy, color='b')
        closeby = zip(*args)
        closeby = filter(lambda x:x[2]<=10, closeby)
        cx,cy = zip(*closeby)[0:2]
        axes.scatter(cx, cy, color='r')
    elif outfile=='dsol_ele_csa_hist':
        ratio,dist = [],[]
        q1,q2,q3,q4 = [],[],[],[]
        for x,y,z in zip(*args):
            if x>=6:
                if -.5<=(y/x)<=.5:
                    ratio.append(y/x)
                    dist.append(z)
            if y>=0 and x>=6:
                q1.append(z)
            elif y>=0 and x<6:
                q2.append(z)
            elif y<0 and x<6:
                q3.append(z)
            elif y<0 and x>=6:
                q4.append(z)
        for q in [q1,q2,q3,q4]:
            print np.average(q)
        axes.hist(dist, bins=20, range=(0,50))
    elif outfile=='ddpk_mddpk_scatter':
        x,y = args[0],args[1]
        axes.scatter(x, y, color='b')
        axes.set_xlim(-20,15)
        axes.set_ylim(-20,15)
    else:
        x,y = args[0],args[1]
        axes.scatter(x, y, color='b')
    pfit = Polyfit(x,y)
    fx = [-20,15]
    fy = []
    for ifx in fx:
        fy.append(pfit['poly'][0]*ifx+pfit['poly'][1])
    axes.plot(fx,fy)
    print pfit['det']
    axes.text(.05,.95, 'm = %3.3f, b = %3.3f, R^2 = %3.3f' % (pfit['poly'][0],pfit['poly'][1],pfit['det']),horizontalalignment='left',verticalalignment='top',transform=axes.transAxes)
    
    fig.savefig(os.path.join(outdir, outfile+'.png'), bbox_inches='tight', transparent=True)
    fig.clf()

def MultiScatter(*args, **kwargs):
    if 'outfile' in kwargs:
        outfile = kwargs['outfile']
    else:
        outfile = 'MultiScatter'
    if 'outdir' in kwargs:
        outdir = kwargs['outdir']
    else:
        outdir = '/home/tel/benchmark/links'
    if 'repickle' in kwargs:
        repickle = kwargs['repickle']
    else:
        repickle = True

    if repickle:
        Repickle(outdir, outfile, *args)
    
    plt.figure(figsize=(10*len(args)/2,10))
    gs = gridspec.GridSpec(1,len(args)/2)
    gs.update(wspace=0.04, hspace=0.07)
    axes = []
    for i,(x,y) in enumerate(zip(args[::2],args[1::2])):
        ax = plt.subplot(gs[0, i])
        axes.append(ax)
        ax.scatter(x, y, color='b')
        pfit = Polyfit(x,y)
        fx = [-100,100]
        fy = []
        for ifx in fx:
            fy.append(pfit['poly'][0]*ifx+pfit['poly'][1])
        ax.plot(fx,fy)
        print pfit['det']
        ax.text(.05,.95, 'm = %3.3f, b = %3.3f, R^2 = %3.3f' % (pfit['poly'][0],pfit['poly'][1],pfit['det']),horizontalalignment='left',verticalalignment='top',transform=ax.transAxes)
        
    for ax in axes:
        for tick in ax.xaxis.get_major_ticks():
            tick.label.set_fontsize('small')
            tick.label.set_rotation(45)
        for tick in ax.yaxis.get_major_ticks():
            tick.label.set_fontsize('small')
    
    d,delta= None,None
    if outfile=='quick_dpk_full_dpk_multiscatter':
        d = 'd'
        delta = r'\Delta '
    elif outfile=='quick_ddpk_full_ddpk_multiscatter':
        d = 'dd'
        delta = r'\Delta\Delta '
    if d!=None and delta!=None:
        xlabels = [r'Quick $%spK_{a}$' % delta]*3
        ylabels = [r'Full $%spK_{a}$' % delta]*3
        titles = ['low dsol', 'mid dsol', 'high dsol']
        for i, ax in enumerate(axes):
            ax.set_xlabel(xlabels[i])
            ax.set_ylabel(ylabels[i])
            ax.set_title(titles[i])
            ax.set_xlim(-15,25)
            ax.set_ylim(-15,25)
            ax.plot([-15,25],[-15,25])
    elem=None
    
    if outfile=='e4_dsol_e20_dsol_multiscatter':
        elem='dsol'
        xlims = ((-5,5),(0,10),(5,15))
        ylims = ((-5,5),(-5,5),(-5,5))
    elif outfile=='e4_ele_e20_ele_multiscatter':
        elem='ele'
        xlims = ((-15,10),(-20,5),(-25,0))
        ylims = ((-20,5),(-20,5),(-20,5))
    elif outfile=='e4_ratio_e20_ratio_multiscatter':
        elem='ratio'
        xlims = ((-10,5),(-15,5),(-10,5))
        ylims = ((-10,5),(-15,5),(-10,5))
    if elem!= None:
        xlabels = [r'%s ${\epsilon}_{p}=4$' % elem]*3
        ylabels = [r'%s ${\epsilon}_{p}=20$' % elem]*3
        titles = ['low dsol', 'mid dsol', 'high dsol']
        for i, ax in enumerate(axes):
            ax.set_xlabel(xlabels[i])
            ax.set_ylabel(ylabels[i])
            ax.set_title(titles[i])
            ax.set_xlim(xlims[i])
            ax.set_ylim(ylims[i])
            ax.plot([-100,100],[-100,100])
    
    plt.savefig(os.path.join(outdir, outfile+'.png'), bbox_inches='tight')#, transparent=True)
    plt.clf()
    
if __name__=="__main__":
    workdir=sys.argv[1]
    #Unpickle(indir=workdir, infile='dsol_ele_scatter', outdir=workdir, outfile='dsol_ele_scatter', argn=3, plotter=Scatter)
    #Unpickle(indir=workdir, infile='dsol_ele_csa_scatter', outdir=workdir, outfile='dsol_ele_csa_scatter', argn=3, plotter=Scatter)
    #Unpickle(indir=workdir, infile='dsol_ele_csa_scatter', outdir=workdir, outfile='dsol_ele_csa_hist', argn=3, plotter=Scatter)
    Unpickle(indir=workdir, infile='ddpk_mddpk_scatter', outdir=workdir, outfile='ddpk_mddpk_scatter', argn=2, plotter=Scatter)
    #Unpickle(indir=workdir, infile='ratio_ddpk_scatter', outdir=workdir, outfile='ratio_ddpk_scatter', argn=2, plotter=Scatter)
    #Unpickle(indir=workdir, infile='quick_dpk_full_dpk_multiscatter', outdir=workdir, outfile='quick_dpk_full_dpk_multiscatter', argn=6, plotter=MultiScatter)
    #Unpickle(indir=workdir, infile='e4_dsol_e20_dsol_multiscatter', outdir=workdir, outfile='e4_dsol_e20_dsol_multiscatter', argn=6, plotter=MultiScatter)
    #Unpickle(indir=workdir, infile='e4_ele_e20_ele_multiscatter', outdir=workdir, outfile='e4_ele_e20_ele_multiscatter', argn=6, plotter=MultiScatter)
    #Unpickle(indir=workdir, infile='e4_ratio_e20_ratio_multiscatter', outdir=workdir, outfile='e4_ratio_e20_ratio_multiscatter', argn=6, plotter=MultiScatter)
    
#    fig = plt.figure(1)
#    axes = plt.axes()
#
#    axes.set_xlabel('Minimum distance\nto active site $(\AA)$', size = 40, multialignment = 'center')
#    axes.set_ylabel(ylabel, size = 40, multialignment = 'center')
#    fig.set_size_inches(11,8)
#    axes.set_xticklabels(['','10','20','30','40','50'])
#    xtick = axes.get_xticklabels()
#    for tick in xtick:
#        tick.set_fontsize(30)
#    #axes.set_yticklabels([])
#    ytick = axes.get_yticklabels()
#    for tick in ytick:
#        tick.set_fontsize(30)
#    axes.hist(x_high_c, bins=30, color='b')
#    axes.set_xlim(0, 50)
#    fig.savefig(fname+'_highc.png', bbox_inches='tight', transparent=True)
#    axes.cla()
#    
#    axes.set_xlabel('Minimum distance\nto active site $(\AA)$', size = 40, multialignment = 'center')
#    axes.set_ylabel(ylabel, size = 40, multialignment = 'center')
#    fig.set_size_inches(11,8)
#    #axes.set_xticklabels([])
#    #axes.set_yticklabels([])
#    axes.set_xticklabels(['','10','20','30','40','50'])
#    xtick = axes.get_xticklabels()
#    for tick in xtick:
#        tick.set_fontsize(30)
#    axes.set_yticklabels(['0','','1','','2'])
#    ytick = axes.get_yticklabels()
#    for tick in ytick:
#        tick.set_fontsize(30)
#    
#    axes.hist(x_low_c, bins=30, color='r')
#    axes.set_xlim(0, 50)
#    fig.savefig(fname+'_lowc.png', bbox_inches='tight', transparent=True)
#    fig.clf()
